// Copyright (C) 2018-2023 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#pragma once

#include "cpu_types.h"
#include "input.h"
#include "memory_state.h"
#include "node.h"
#include "ov_optional.hpp"
#include "proxy_mem_mgr.h"

#include <map>

namespace ov {
namespace intel_cpu {
namespace node {

class MemoryOutputBase;
class MemoryInputBase;
class ScaledDotProductAttention;

class MemoryNode {
 public:
    explicit MemoryNode(std::string id) : m_id(id) {}
    explicit MemoryNode(const std::shared_ptr<ov::Node>& op);
    virtual ~MemoryNode() = default;
    std::string getId() const {
        return m_id;
    }

private:
    std::string m_id;
};

class MemoryStateNode : public MemoryNode {
public:
    using MemoryNode::MemoryNode;
    virtual void assignState(MemStatePtr newState) = 0;
    virtual MemStatePtr makeState() const = 0;
};

/**
 * @brief
 * TODO: ATTENTION: this is a temporary solution, this connection should be keep in graph
 * WARNING: thread_local and holderMutex are not needed if moved into graph
 */
class MemoryNodeVirtualEdge {
public:
    using Holder = std::map<std::string, MemoryNode*>;
    static Holder & getExisted() {
        thread_local static Holder existed;
        return existed;
    }

    static MemoryNode * getByName(Holder& holder, std::string name) {
        auto result = holder.find(name);
        if (result != holder.end()) {
            return result->second;
        }
        return nullptr;
    }

    static Holder* registerOutput(MemoryOutputBase * node);
    static Holder* registerInput(MemoryInputBase * node);
    static void remove(MemoryNode * node, Holder* holder);
    static std::mutex holderMutex;
};

class MemoryOutputBase : public Node, public MemoryNode {
public:
    MemoryOutputBase(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context);
    MemoryOutputBase(const std::string id,
                     const std::string& name,
                     const std::string& type,
                     const Shape& input_shape,
                     const ov::element::Type& input_prc,
                     const GraphContext::CPtr context);

    ~MemoryOutputBase() override;
    static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;
    void getSupportedDescriptors() override;
    void initSupportedPrimitiveDescriptors() override;
    void initOptimalPrimitiveDescriptor() override;
    void createPrimitive() override {}
    bool created() const override {
        return getType() == Type::MemoryOutput;
    }

    void registerInputNode(MemoryInputBase* node);
    void deregisterSibling(MemoryInputBase* node);

    bool needShapeInfer() const override { return false; }
    bool needPrepareParams() const override { return false; }

    virtual void assignExtMemory(const MemoryPtr& mem, const MemoryDescPtr& memDesc) = 0;

protected:
    MemoryInputBase& getInputNode();

private:
    /**
     * @brief keeps reference to input sibling node
     */
    MemoryInputBase* inputNode = nullptr;
    MemoryNodeVirtualEdge::Holder* holder = nullptr;
};

class MemoryOutput : public MemoryOutputBase {
public:
    using MemoryOutputBase::MemoryOutputBase;
    static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;

    void execute(dnnl::stream strm) override;
    void executeDynamicImpl(dnnl::stream strm) override;

    void resolveInPlaceEdges(Edge::LOOK look) override;

    void assignExtMemory(const MemoryPtr& mem, const MemoryDescPtr& memDesc) override;

private:
    MemoryPtr assignedMem = nullptr;
    MemoryDescPtr extMemDesc = nullptr; // used for resize
    ProxyMemoryMngrPtr memMngr = nullptr;
};

class MemoryOutputStub : public MemoryOutputBase {
public:
    using MemoryOutputBase::MemoryOutputBase;
    static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;

    void execute(dnnl::stream strm) override;
    void executeDynamicImpl(dnnl::stream strm) override;

    void resolveInPlaceEdges(Edge::LOOK look) override;

    void assignExtMemory(const MemoryPtr& mem, const MemoryDescPtr& memDesc) override;
    bool isExecutable() const override;
};

class MemoryInputBase : public Input, public MemoryStateNode {
public:
    MemoryInputBase(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context);
    MemoryInputBase(const std::string id,
                    const std::string& name,
                    const std::string& type,
                    const Shape& output_shape,
                    const ov::element::Type& output_prc,
                    const GraphContext::CPtr context,
                    const ov::optional<Shape>& input_shape,
                    const ov::optional<ov::element::Type>& input_prc);

    ~MemoryInputBase() override;

    static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;
    bool created() const override {
        return getType() == Type::MemoryInput;
    }

    void initSupportedPrimitiveDescriptors() override;

    void registerOutputNode(MemoryOutputBase* node);
    void deregisterSibling(MemoryOutputBase* node);

    MemoryOutputBase& getOutputNode();

private:
    /**
     * @brief keeps reference to output sibling node
     */
    MemoryOutputBase* outputNode = nullptr;
    MemoryNodeVirtualEdge::Holder* holder = nullptr;
};

class MemoryInput : public MemoryInputBase {
public:
    using MemoryInputBase::MemoryInputBase;
    static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;

    bool needShapeInfer() const override;
    bool isExecutable() const override;
    void initOptimalPrimitiveDescriptor() override;
    void execute(dnnl::stream strm) override;
    void executeDynamicImpl(dnnl::stream strm) override;

    void resolveInPlaceEdges(Edge::LOOK look) override;

    void assignState(MemStatePtr newState) override;
    MemStatePtr makeState() const override;

private:
    bool isExecutableFlag = true;
    ProxyMemoryMngrPtr memMngr = nullptr;
    MemoryPtr assignedMem = nullptr;
};

class MemoryInputSDPA : public MemoryInputBase {
public:
    MemoryInputSDPA(const std::string id,
                    const std::string& name,
                    const std::string& type,
                    const Shape& output_shape,
                    const ov::element::Type& output_prc,
                    const GraphContext::CPtr context,
                    const ov::optional<Shape>& input_shape,
                    const ov::optional<ov::element::Type>& input_prc,
                    const std::shared_ptr<ScaledDotProductAttention>& sdpaNode);

    static bool isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std::string& errorMessage) noexcept;

    bool needShapeInfer() const override;
    bool isExecutable() const override;

    void createPrimitive() override;
    void initSupportedPrimitiveDescriptors() override;
    void initOptimalPrimitiveDescriptor() override;

    void execute(dnnl::stream strm) override;
    void executeDynamicImpl(dnnl::stream strm) override;

    void resolveInPlaceEdges(Edge::LOOK look) override;

    void assignState(MemStatePtr newState) override;
    MemStatePtr makeState() const override;

private:
    std::weak_ptr<ScaledDotProductAttention> m_sdpaNode;
    int m_child_port_idx = -1;
    bool m_needShapeInfer = false;
};
}   // namespace node
}   // namespace intel_cpu
}   // namespace ov
